# 团队状态保存和加载

from autogen_agentchat.agents import AssistantAgent
from autogen_agentchat.base import TaskResult
from autogen_agentchat.conditions import TextMentionTermination, ExternalTermination, MaxMessageTermination, \
    TextMessageTermination
from autogen_agentchat.messages import TextMessage
from autogen_agentchat.teams import RoundRobinGroupChat
from autogen_agentchat.ui import Console
from autogen_core import CancellationToken
from autogen_ext.models.openai import OpenAIChatCompletionClient

model_client = OpenAIChatCompletionClient(model="qwen2.5:7b",
                                          model_info={
                                              "vision": False,
                                              "function_calling": True,
                                              "family": "Qwen3",
                                              "structured_output": True,
                                              "json_output": True,
                                          },
                                          api_key="ollama",
                                          base_url="http://127.0.0.1:11434/v1")

async def main():

    # 创建一个使用 increment_number 函数的工具代理。
    looped_assistant = AssistantAgent(
        "looped_assistant",
        model_client=model_client,
        system_message="你是一个帮助AI助手，善于写诗。",
    )
    # 当代理以文本消息响应时停止任务的终止条件。
    termination_condition = MaxMessageTermination(max_messages=3)
    # 创建一个由循环助手代理和终止条件组成的团队。
    team = RoundRobinGroupChat(
        [looped_assistant],
        termination_condition=termination_condition,
    )
    # 使用一个任务运行团队，并将消息打印到控制台。
    await Console(team.run_stream(task="写一首关于春天的绝句."))
    # 保存团队状态
    team_state = await team.save_state()
    # 重置团队
    await team.reset()
    print("-"*50)
    print("重置团队后，继续任务")
    await Console(team.run_stream(task="最后一句是什么?"))
    # 加载团队状态
    print("-"*50)
    print("加载团队状态")
    
    await team.load_state(team_state)
    # 继续任务
    await Console(team.run_stream(task="最后一句是什么?"))
if __name__ == "__main__":
    import asyncio
    asyncio.run(main())